/*
 * Copyright (C) 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.android.tools.profilers.perfetto.traceprocessor

import com.android.tools.profiler.perfetto.proto.TraceProcessor
import com.android.tools.profiler.perfetto.proto.TraceProcessor.AndroidFrameEventsResult.*
import com.android.tools.profiler.perfetto.proto.TraceProcessor.PowerCounterTracksResult
import com.android.tools.profilers.cpu.ThreadState
import com.android.tools.profilers.cpu.systemtrace.AndroidFrameTimelineEvent
import com.android.tools.profilers.cpu.systemtrace.CounterModel
import com.android.tools.profilers.cpu.systemtrace.ThreadStateModel
import com.android.tools.profilers.cpu.systemtrace.CpuCoreModel
import com.android.tools.profilers.cpu.systemtrace.ProcessModel
import com.android.tools.profilers.cpu.systemtrace.SchedulingEventModel
import com.android.tools.profilers.cpu.systemtrace.SystemTraceModelAdapter
import com.android.tools.profilers.cpu.systemtrace.ThreadModel
import com.android.tools.profilers.cpu.systemtrace.TraceEventModel
import perfetto.protos.PerfettoTrace
import java.io.Serializable
import java.util.Deque
import java.util.LinkedList
import java.util.concurrent.TimeUnit
import com.android.tools.profilers.cpu.config.ProfilingConfiguration.TraceType
import com.android.tools.profilers.perfetto.traceprocessor.TraceProcessorModelUtils.findValueNearKey
import org.jetbrains.annotations.VisibleForTesting
import kotlin.math.max
import kotlin.math.min

class TraceProcessorModel(builder: Builder) : SystemTraceModelAdapter, Serializable {

  companion object {
    // generated by serialver
    @JvmStatic
    val serialVersionUID = -2228053132990163476L
  }

  private val processMap: Map<Int, ProcessModel>
  private val cpuCores: List<CpuCoreModel>
  private val powerRails: List<CounterModel>
  private val batteryDrain: List<CounterModel>
  private val androidFrameLayers: List<Layer>
  private val androidFrameTimelineEvents: List<AndroidFrameTimelineEvent>

  private val danglingThreads = builder.danglingThreads

  private val startCaptureTimestamp = builder.startCaptureTimestamp
  private val endCaptureTimestamp = builder.endCaptureTimestamp

  init {
    // Build processMap
    val processMapBuilder = mutableMapOf<Int, ProcessModel>()
    for (process in builder.processById.values) {
      val updatedThreadMap = process.threadById.mapValues { entry ->
        entry.value.copy(
          traceEvents = builder.threadToEventsMap.getOrDefault(entry.key, listOf()),
          schedulingEvents = builder.threadToScheduling.getOrDefault(entry.key, listOf()),
          threadStateEvents = builder.threadToThreadStates.getOrDefault(entry.key, listOf())
        )
      }.toSortedMap()

      val counterMap = builder.processToCounters.getOrDefault(process.id, listOf())
        .associateBy { it.name }
      processMapBuilder[process.id] = process.copy(threadById = updatedThreadMap, counterByName = counterMap)
    }
    processMap = processMapBuilder.toSortedMap()

    // Power counter data contains two types of counters: power rail and battery drain.
    // The name of each power rail is prefixed by 'power.' while
    // the name of each battery drain metric is prefixed by 'batt.'.
    powerRails = builder.powerCounters.filter { it.name.contains("power.") }
    batteryDrain = builder.powerCounters.filter { it.name.contains("batt.") }

    // Build cpuCores
    cpuCores = (0 until builder.cpuCoresCount).map {
      val cpuCountersMap = builder.coreToCpuCounters.getOrDefault(it, listOf())
        .associateBy { counter -> counter.name }
      CpuCoreModel(it, builder.coreToScheduling.getOrDefault(it, listOf()), cpuCountersMap)
    }
    androidFrameTimelineEvents = builder.androidFrameTimelineEvents
    androidFrameLayers = when {
      androidFrameTimelineEvents.isEmpty() -> builder.androidFrameLayers
      else -> builder.androidFrameLayers.renumbered(androidFrameTimelineEvents, builder.surfaceflingerDisplayTokenToEndNs)
    }
  }

  override fun getCaptureStartTimestampUs() = startCaptureTimestamp
  override fun getCaptureEndTimestampUs() = endCaptureTimestamp

  override fun getProcessById(id: Int) = processMap[id]
  override fun getProcesses() = processMap.values.toList()
  override fun getDanglingThread(tid: Int): ThreadModel? = danglingThreads[tid]
  override fun getCpuCores() = cpuCores
  override fun getPowerRails(): List<CounterModel> = powerRails
  override fun getBatteryDrain(): List<CounterModel> = batteryDrain
  override fun getSystemTraceTechnology() = TraceType.PERFETTO

  // TODO(b/156578844): Fetch data from TraceProcessor error table to populate this.
  override fun isCapturePossibleCorrupted() = false
  override fun getAndroidFrameLayers() = androidFrameLayers
  override fun getAndroidFrameTimelineEvents() = androidFrameTimelineEvents

  class Builder {
    internal var startCaptureTimestamp = Long.MAX_VALUE
    internal var endCaptureTimestamp = Long.MIN_VALUE
    internal var cpuCoresCount = 0
    internal val processById = mutableMapOf<Int, ProcessModel>()
    internal val danglingThreads = mutableMapOf<Int, ThreadModel>()
    internal val threadToEventsMap = mutableMapOf<Int, List<TraceEventModel>>()
    internal val threadToScheduling = mutableMapOf<Int, List<SchedulingEventModel>>()
    internal val coreToScheduling = mutableMapOf<Int, List<SchedulingEventModel>>()
    internal val coreToCpuCounters = mutableMapOf<Int, List<CounterModel>>()
    internal val processToCounters = mutableMapOf<Int, List<CounterModel>>()
    internal val threadToThreadStates = mutableMapOf<Int, List<ThreadStateModel>>()
    internal val powerCounters = mutableListOf<CounterModel>()
    internal val androidFrameLayers = mutableListOf<Layer>()
    internal val androidFrameTimelineEvents = mutableListOf<AndroidFrameTimelineEvent>()
    internal var surfaceflingerDisplayTokenToEndNs = mapOf<Long, Long>()

    companion object {
      /**
       * Returns the sum of the durations of running scheduling events that are within the range between `rangeStartUs` and `rangeEndUs`.
       * Assumes the given `schedulingEvents` is already sorted by `startTimestampUs`.
       */
      @VisibleForTesting
      fun getRunningTimeUsInRange(rangeStartUs: Long, rangeEndUs: Long, schedulingEvents: List<SchedulingEventModel>): Long {
        var firstSchedIndexToCheck = schedulingEvents.binarySearch { sched -> sched.endTimestampUs.compareTo(rangeStartUs) }
        if (firstSchedIndexToCheck < 0)
          firstSchedIndexToCheck = -(firstSchedIndexToCheck + 1)
        var lastSchedIndexToCheck = schedulingEvents.binarySearch { sched -> sched.startTimestampUs.compareTo(rangeEndUs) }
        if (lastSchedIndexToCheck < 0)
          lastSchedIndexToCheck = -(lastSchedIndexToCheck + 1)
        return schedulingEvents.subList(firstSchedIndexToCheck /* inclusive */,
                                        min(lastSchedIndexToCheck + 1, schedulingEvents.size) /* exclusive */)
          .filter { it.state in setOf(ThreadState.RUNNING, ThreadState.RUNNING_CAPTURED) }
          .filter { it.endTimestampUs >= rangeStartUs && it.startTimestampUs <= rangeEndUs }
          .sumOf { min(it.endTimestampUs, rangeEndUs) - max(it.startTimestampUs, rangeStartUs) }
      }
    }

    fun addProcessMetadata(processMetadataResult: TraceProcessor.ProcessMetadataResult) {
      for (process in processMetadataResult.processList) {
        processById[process.id.toInt()] = ProcessModel(
          process.id.toInt(),
          process.name,
          process.threadList.associate {t -> t.id.toInt() to ThreadModel(t.id.toInt(),
                                                                         process.id.toInt(),
                                                                         t.name,
                                                                         listOf(),
                                                                         listOf(),
                                                                         listOf()) }.toSortedMap(),
          mapOf())
      }

      for (thread in processMetadataResult.danglingThreadList) {
        danglingThreads[thread.id.toInt()] = ThreadModel(thread.id.toInt(), 0, thread.name, emptyList(), emptyList(), emptyList())
      }
    }

    fun addTraceEvents(traceEventsResult: TraceProcessor.TraceEventsResult) {
      for (thread in traceEventsResult.threadList) {
        val rootIds = mutableSetOf<Long>()
        val eventToChildrenIds = mutableMapOf<Long, MutableList<Long>>()
        val eventPerId = mutableMapOf<Long, TraceEventModel>()
        val incompleteEventToThreadId = mutableMapOf<Long, Int>()

        for (event in thread.traceEventList) {
          if (event.depth > 0) {
            eventToChildrenIds.getOrPut(event.parentId) { mutableListOf() }.add(event.id)
          }
          else {
            rootIds.add(event.id)
          }

          val startTimestampUs = convertToUs(event.timestampNanoseconds)
          var endTimestampUs = -1L
          var cpuTimeUs = -1L
          if (event.durationNanoseconds == -1L){
            // A duration of -1 signifies an incomplete trace event. For these events, CPU time calculation
            // is deferred until the end time is estimated. Thread scheduling information is stored to facilitate
            // CPU time calculation for such events at a later stage
            incompleteEventToThreadId[event.id] = thread.threadId.toInt()
          }
          else {
            val durationTimestampUs = convertToUs(event.durationNanoseconds)
            endTimestampUs = startTimestampUs + durationTimestampUs
            cpuTimeUs = getRunningTimeUsInRange(startTimestampUs, endTimestampUs,
                                                threadToScheduling.getOrDefault(thread.threadId.toInt(), listOf()))
          }

          eventPerId[event.id] = TraceEventModel(event.name,
                                                 startTimestampUs,
                                                 endTimestampUs,
                                                 cpuTimeUs,
                                                 listOf())
        }

        val reconstructedTree = reconstructTraceTree(rootIds, eventToChildrenIds, eventPerId, incompleteEventToThreadId)
        threadToEventsMap[thread.threadId.toInt()] = rootIds.mapNotNull { reconstructedTree[it] }
      }
    }

    // Runs through the partially computed events to rebuild the whole trace trees, by doing a DFS from the root nodes.
    private fun reconstructTraceTree(
      rootIds: Set<Long>, eventToChildrenIds: Map<Long, List<Long>>, eventPerId: Map<Long, TraceEventModel>, incompleteEventToThreadId:  Map<Long, Int>): Map<Long, TraceEventModel> {

      val reconstructedEventsPerId = mutableMapOf<Long, TraceEventModel>()

      val visitedAllChildren = mutableSetOf<Long>()
      val eventIdStack: Deque<Long> = LinkedList(rootIds)

      while (eventIdStack.isNotEmpty()) {
        val eventId = eventIdStack.first

        // If we have not visited this node yet, then we need to push all its children to the front of the queue
        // and continue the main loop. Next time we pass on this one, we will process it as we know all its children
        // have been processed already.
        if (!visitedAllChildren.contains(eventId)) {
          eventToChildrenIds.getOrDefault(eventId, mutableListOf()).forEach { eventIdStack.addFirst(it) }
          visitedAllChildren.add(eventId)
          continue
        }

        eventIdStack.removeFirst()
        val children = eventToChildrenIds.getOrDefault(eventId, mutableListOf())
          .map { reconstructedEventsPerId[it] ?: error("Children should have been computed already") }
          .sortedBy { it.startTimestampUs }

        val event = eventPerId[eventId] ?: error("Trace Event should be present in the map")

        val myStart = event.startTimestampUs
        val maxEndTs = children.lastOrNull()?.endTimestampUs ?: 0L
        // For incomplete events, the end time is estimated as the maximum of its start time and its last child's end time.
        // The CPU time is calculated using the estimated end time
        val updatedEndTimestampUs = if (event.endTimestampUs != -1L) event.endTimestampUs else maxOf(myStart, maxEndTs)
        val updatedCpuTimeUs = if (event.cpuTimeUs != -1L) event.cpuTimeUs else getRunningTimeUsInRange(event.startTimestampUs, updatedEndTimestampUs,
                                                                                                        threadToScheduling.getOrDefault(incompleteEventToThreadId[eventId], listOf()))
        val updatedEvent = event.copy(
          endTimestampUs = updatedEndTimestampUs,
          cpuTimeUs = updatedCpuTimeUs,
          childrenEvents = children)
        reconstructedEventsPerId[eventId] = updatedEvent

        // Update the global start/end of the capture.
        startCaptureTimestamp = minOf(startCaptureTimestamp, updatedEvent.startTimestampUs)
        endCaptureTimestamp = maxOf(endCaptureTimestamp, updatedEvent.endTimestampUs)
      }

      return reconstructedEventsPerId
    }

    fun addSchedulingEvents(schedEvents: TraceProcessor.SchedulingEventsResult) {
      cpuCoresCount = maxOf(cpuCoresCount, schedEvents.numCores)

      val perThreadScheduling = mutableMapOf<Int, MutableList<SchedulingEventModel>>()
      val perCoreScheduling = mutableMapOf<Int, MutableList<SchedulingEventModel>>()
      schedEvents.schedEventList
        .groupBy { it.threadId }
        .forEach { (tid, events) ->
          events.forEachIndexed { index, event ->
            val startTimestampUs = convertToUs(event.timestampNanoseconds)
            val durationUs = convertToUs(event.durationNanoseconds)
            val endTimestampUs = startTimestampUs + durationUs
            startCaptureTimestamp = minOf(startCaptureTimestamp, startTimestampUs)
            endCaptureTimestamp = maxOf(endCaptureTimestamp, endTimestampUs)

            // Scheduling events from the Trace Processor encode thread states differently from ftrace.
            // TP only records the end_state of the event and implies RUNNING as the start_state always:
            // https://perfetto.dev/docs/data-sources/cpu-scheduling#decoding-code-end_state-code-
            //
            // So for every event except the last one, we need to insert a RUNNING event + an end_state event.
            val schedEvent = SchedulingEventModel(ThreadState.RUNNING_CAPTURED,
                                                  startTimestampUs,
                                                  endTimestampUs,
                                                  durationUs,
                                                  // Parameter `cpuTimeUs` is not used, so just use `durationUs` as a placeholder
                                                  durationUs,
                                                  event.processId.toInt(),
                                                  event.threadId.toInt(),
                                                  event.cpu)
            // Add a RUNNING event and an [end_state] event to thread scheduling events.
            perThreadScheduling.getOrPut(tid.toInt()) { mutableListOf() }.apply {
              // The RUNNING thread state event.
              add(schedEvent)
              if (index < events.size - 1) {
                val nextStartTimestampUs = convertToUs(events[index + 1].timestampNanoseconds)
                val nextDurationTimeUs = nextStartTimestampUs - endTimestampUs
                // The [end_state] thread state event.
                add(SchedulingEventModel(convertSchedulingState(event.endState),
                                         endTimestampUs,
                                         nextStartTimestampUs,
                                         nextDurationTimeUs,
                                         nextDurationTimeUs,
                                         event.processId.toInt(),
                                         event.threadId.toInt(),
                                         event.cpu))
              }
            }
            // Add just the RUNNING event to core scheduling events.
            perCoreScheduling.getOrPut(event.cpu) { mutableListOf() }.add(schedEvent)
          }
        }

      perThreadScheduling.forEach {
        val previousList = threadToScheduling[it.key] ?: listOf()
        threadToScheduling[it.key] = previousList.plus(it.value).sortedBy { s -> s.startTimestampUs }
      }
      perCoreScheduling.forEach {
        val previousList = coreToScheduling[it.key] ?: listOf()
        coreToScheduling[it.key] = previousList.plus(it.value).sortedBy { s -> s.startTimestampUs }
      }
    }

    private fun convertSchedulingState(state: TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState): ThreadState {
      return when (state) {
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.RUNNABLE -> ThreadState.RUNNABLE_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.RUNNABLE_PREEMPTED -> ThreadState.RUNNABLE_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.DEAD -> ThreadState.DEAD_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.SLEEPING -> ThreadState.SLEEPING_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.SLEEPING_UNINTERRUPTIBLE -> ThreadState.WAITING_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.WAKING -> ThreadState.RUNNABLE_CAPTURED
        TraceProcessor.SchedulingEventsResult.SchedulingEvent.SchedulingState.WAKE_KILL -> ThreadState.RUNNABLE_CAPTURED
        else -> ThreadState.UNKNOWN
      }
    }

    fun addCpuCounters(result: TraceProcessor.CpuCoreCountersResult) {
      cpuCoresCount = maxOf(cpuCoresCount, result.numCores)

      result.countersPerCoreList.forEach { countersPerCore ->
        coreToCpuCounters[countersPerCore.cpu] = countersPerCore.counterList.map { counter ->
          CounterModel(counter.name, counter.valueList.associate {
            convertToUs(it.timestampNanoseconds) to it.value
          }.toSortedMap())
        }
      }
    }

    fun addProcessCounters(counters: TraceProcessor.ProcessCountersResult) {
      processToCounters[counters.processId.toInt()] = counters.counterList.map { counter ->
        CounterModel(counter.name, counter.valueList.associate {
          convertToUs(it.timestampNanoseconds) to it.value
        }.toSortedMap())
      }
    }

    fun addThreadStates(states: TraceProcessor.ThreadStatesResult) {
      val perThreadState = mutableMapOf<Int, MutableList<ThreadStateModel>>()

      states.stateEventList.groupBy { it.threadId }.forEach { (tid, events) ->
        events.forEach {
          val threadStateEvent =
            ThreadStateModel(if (it.state.running) ThreadState.RUNNING_CAPTURED else convertSchedulingState(it.state.nonRunning),
                             convertToUs(it.timestampNanoseconds),
                             convertToUs(it.timestampNanoseconds + it.durationNanoseconds))
          perThreadState.getOrPut(tid.toInt()) { mutableListOf() }.apply {
            add(threadStateEvent)
          }
        }
      }

      perThreadState.forEach {
        val previousList = threadToThreadStates[it.key] ?: listOf()
        threadToThreadStates[it.key] = previousList.plus(it.value).sortedBy { s -> s.startTimestampUs }
      }
    }

    fun addPowerCounters(counters: PowerCounterTracksResult) {
      // powerCounters include both power rail counter and battery drain counter data
      powerCounters.addAll(counters.counterList.map { counter ->
        CounterModel(counter.name, counter.valueList.associate {
          convertToUs(it.timestampNanoseconds) to it.value
        }.toSortedMap())
      })
    }

    fun addAndroidFrameEvents(frameEventsResult: TraceProcessor.AndroidFrameEventsResult) {
      androidFrameLayers.addAll(frameEventsResult.layerList)
    }

    /**
     * Process the Frame Timeline query result and turn it into [AndroidFrameTimelineEvent]s, ordered by
     * [AndroidFrameTimelineEvent.expectedStartUs]
     */
    fun addAndroidFrameTimelineEvents(frameTimelineResult: TraceProcessor.AndroidFrameTimelineResult) {
      // Match actual timeline slices to expected timeline slices by surfaceFrameToken and construct an AndroidFrameTimelineEvent
      // from two matching slices.
      val expectedSlicesBySurfaceFrame = frameTimelineResult.expectedSliceList.associateBy { it.surfaceFrameToken }
      frameTimelineResult.actualSliceList.forEach { actualSlice ->
        expectedSlicesBySurfaceFrame[actualSlice.surfaceFrameToken]?.let { expectedSlice ->
          val expectedEndUs = convertToUs(expectedSlice.timestampNanoseconds + expectedSlice.durationNanoseconds)
          val actualEndUs = convertToUs(actualSlice.timestampNanoseconds + actualSlice.durationNanoseconds)
          androidFrameTimelineEvents.add(AndroidFrameTimelineEvent(expectedSlice.displayFrameToken,
                                                                   expectedSlice.surfaceFrameToken,
                                                                   convertToUs(expectedSlice.timestampNanoseconds),
                                                                   expectedEndUs,
                                                                   actualEndUs,
                                                                   actualSlice.layerName,
                                                                   parsePresentType(actualSlice.presentType),
                                                                   parseAppJankType(actualSlice.jankType),
                                                                   actualSlice.onTimeFinish,
                                                                   actualSlice.gpuComposition,
                                                                   actualSlice.layoutDepth))
        }
      }
      androidFrameTimelineEvents.sortBy { it.expectedStartUs }
    }

    fun indexSurfaceflingerFrameTimelineEvents(frameTimelineResult: TraceProcessor.AndroidFrameTimelineResult) {
      surfaceflingerDisplayTokenToEndNs = frameTimelineResult.actualSliceList.associate {
        it.displayFrameToken to it.timestampNanoseconds + it.durationNanoseconds
      }
    }

    fun build(): TraceProcessorModel {
      return TraceProcessorModel(this)
    }

    private fun convertToUs(tsNanos: Long) = TimeUnit.NANOSECONDS.toMicros(tsNanos)

    private fun parsePresentType(presetTypeStr: String) = when (presetTypeStr) {
      "On-time Present" -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_ON_TIME
      "Late Present" -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_LATE
      "Early Present" -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_EARLY
      "Dropped Frame" -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_DROPPED
      "Unknown Present" -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_UNKNOWN
      else -> PerfettoTrace.FrameTimelineEvent.PresentType.PRESENT_UNSPECIFIED
    }

    private fun parseAppJankType(jankTypeStr: String): PerfettoTrace.FrameTimelineEvent.JankType {
      // jankTypeStr is a comma-separated string of both an app jank type and a surfaceflinger jank type (if exists).
      // In SystemTrace we only care about the app jank type, so once a match is found we can return early.
      jankTypeStr.split(", ").forEach {
        when (it) {
          "None" -> return PerfettoTrace.FrameTimelineEvent.JankType.JANK_NONE
          "App Deadline Missed" -> return PerfettoTrace.FrameTimelineEvent.JankType.JANK_APP_DEADLINE_MISSED
          "Buffer Stuffing" -> return PerfettoTrace.FrameTimelineEvent.JankType.JANK_BUFFER_STUFFING
          "Unknown Jank" -> return PerfettoTrace.FrameTimelineEvent.JankType.JANK_UNKNOWN
        }
      }
      return PerfettoTrace.FrameTimelineEvent.JankType.JANK_UNSPECIFIED
    }
  }
}

private fun List<Layer>.renumbered(timelineEvents: List<AndroidFrameTimelineEvent>,
                                   surfaceflingerDisplayTokenToEndNs: Map<Long, Long>): List<Layer> =
  mapFrameNumber(this, timelineEvents, surfaceflingerDisplayTokenToEndNs).let { frameNumberMap ->
    fun timelineNumberOf(lifecycleNumber: Int) = frameNumberMap[lifecycleNumber]?.toInt()
    fun transformFrame(evt: FrameEvent): FrameEvent? =
      timelineNumberOf(evt.frameNumber)?.let { evt.toBuilder().setFrameNumber(it).build() }
    fun transformPhase(phase: Phase): Phase =
      phase.toBuilder().clearFrameEvent().addAllFrameEvent(phase.frameEventList.mapNotNull(::transformFrame)).build()
    fun transformLayer(layer: Layer): Layer =
      layer.toBuilder().clearPhase().addAllPhase(layer.phaseList.mapNotNull(::transformPhase)).build()
    map(::transformLayer)
  }

fun List<Layer>.groupedByPhase(): List<Phase> = asSequence()
  .flatMap(Layer::getPhaseList)
  .groupingBy(Phase::getPhaseName)
  // for each phase, accumulate event list and max depth
  .fold({ _, _ -> mutableListOf<FrameEvent>() to 0}) { phaseName, (eventsAcc, depthAcc), phase -> when (phaseName) {
    "Display" -> phase.frameEventList to 0 // "Display" track is special, so no accumulation
    else -> {
      eventsAcc.addAll(phase.frameEventList.map { e -> e.toBuilder().setDepth(depthAcc + e.depth).build() })
      eventsAcc to (depthAcc + (phase.frameEventList.maxOfOrNull(FrameEvent::getDepth)?.let(Int::inc) ?: 0))
    }
  } }
  .map { (name, result) -> Phase.newBuilder().setPhaseName(name).addAllFrameEvent(result.first).build() }

/**
 * Map lifecycle's frame numbers to timeline's frame numbers by:
 * - identifying the surfaceflinger event sharing the same display-token as the app event
 * - looking up the lifecycle's event whose start matches the surfaceflinger event's end
 */
private fun mapFrameNumber(layers: List<Layer>,
                           timelineEvents: List<AndroidFrameTimelineEvent>,
                           surfaceflingerDisplayTokenToEndNs: Map<Long, Long>): Map<Int, Long> {
  val lifeCycleEventStartNsToNumber = LinkedHashMap<Long, Int>()
  layers.forEach { layer ->
    layer.phaseList.filter {
      // The "Display" phase contains the entire lifecycle of the frame and thus should be the only phase used to match timestamps.
      phase -> phase.phaseName == "Display"
    }.forEach { phase ->
      phase.frameEventList.forEach { event -> lifeCycleEventStartNsToNumber[event.timestampNanoseconds] = event.frameNumber }
    }
  }

  fun getLifecycleFrameNumber(event: AndroidFrameTimelineEvent) : Int? {
    val surfaceflingerEndNs = surfaceflingerDisplayTokenToEndNs[event.displayFrameToken]
    // A tolerance (1000 ns) is given with timestamp-based matching to account for different clocks used for the different data sources.
    // The value of 1000 ns is reasonable given a frame display duration will be at least in the millions of nanoseconds (milliseconds).
    return surfaceflingerEndNs?.let { findValueNearKey(sortedMap = lifeCycleEventStartNsToNumber, targetKey = it, tolerance = 1000L) }
  }

  return timelineEvents
    .filter { event -> getLifecycleFrameNumber(event) != null }
    .associate { event -> getLifecycleFrameNumber(event)!! to event.surfaceFrameToken }
}